#! /usr/bin/env python3
from pyspark.ml import PipelineModel
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

from gai.v2.spark.feature.variable_stats import prepare_test_dataframe_dup_name
from gai.v2.spark.transformer.column_mapper import ColumnMapper


def _make_add_prefix(prefix: str, separator=","):
    @udf(returnType=StringType())
    def add_prefix(cell_text):
        if len(cell_text) != 0:  # note that ``"".split(",")`` returns ``['']``
            prefixed_parts = [prefix + x for x in cell_text.split(separator)]
            cell_text = separator.join(prefixed_parts)
        return cell_text

    return add_prefix


class PrefixAdder(ColumnMapper):
    """Given the name of the input column, the name of the output column,
    the prefix string, and the separator string, an instance of ``PrefixAdder``
    can split each cell of input column with the separator string, and prepend
    the prefix string to each element of the result, and finally combine the
    prefixed elements back to a string with the same separator, and append the
    combined string as an output column.

    Args:
        inputCol: the name of input column
        outputCol: the name of output column
        prefix: the prefix string
        separator: the string used to separate the encoded fields in a cell of ``inputCol``

    >>> from gai.v2.spark.transformer import PrefixAdder
    >>> df = prepare_test_dataframe_dup_name()
    >>> df.show()
    +---+------+------------+-------------------+-----+------------+
    |age|  name|    app_list|      category_list|label|install_pkgs|
    +---+------+------------+-------------------+-----+------------+
    |  7| Alice| app_1,app_3|  cate_0:3,cate_1:2|    0|       app_1|
    |  7|   Bob|app_3,app_21| cate_0:4,cate_10:7|    1|app_4#app_21|
    |  3|Claire| app_1,app_7| cate_0:3,cate_11:4|    0|            |
    |  7|   Dan| app_9,app_5|cate_20:2,cate_99:7|    0|            |
    +---+------+------------+-------------------+-----+------------+
    <BLANKLINE>
    >>> prefix_adder_1 = PrefixAdder(inputCol='app_list',
    ...                         outputCol='prefixed_app_list',
    ...                         prefix='all_',
    ...                         separator=',')
    >>> prefix_adder_2 = PrefixAdder(inputCol='install_pkgs',
    ...                         outputCol='prefixed_install_pkgs',
    ...                         prefix='inst_',
    ...                         separator='#')
    >>> pipeline_model = PipelineModel(stages=[prefix_adder_1, prefix_adder_2])
    >>> prefixed_df = pipeline_model.transform(df)
    >>> prefixed_df.show()
    +---+------+------------+-------------------+-----+------------+--------------------+---------------------+
    |age|  name|    app_list|      category_list|label|install_pkgs|   prefixed_app_list|prefixed_install_pkgs|
    +---+------+------------+-------------------+-----+------------+--------------------+---------------------+
    |  7| Alice| app_1,app_3|  cate_0:3,cate_1:2|    0|       app_1| all_app_1,all_app_3|           inst_app_1|
    |  7|   Bob|app_3,app_21| cate_0:4,cate_10:7|    1|app_4#app_21|all_app_3,all_app_21| inst_app_4#inst_a...|
    |  3|Claire| app_1,app_7| cate_0:3,cate_11:4|    0|            | all_app_1,all_app_7|                     |
    |  7|   Dan| app_9,app_5|cate_20:2,cate_99:7|    0|            | all_app_9,all_app_5|                     |
    +---+------+------------+-------------------+-----+------------+--------------------+---------------------+
    <BLANKLINE>
    """

    def __init__(self, inputCol, outputCol, prefix: str, separator=","):
        super(PrefixAdder, self).__init__(
            fun=_make_add_prefix(prefix, separator),
            inputCols=[inputCol],
            outputCol=outputCol
        )
